import base64
import json
import logging
import os
import re
from datetime import datetime

import torch
from PIL import Image
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from paddleocr import PaddleOCR
from torchvision import transforms
import pandas as pd
import cv2
import numpy as np

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 加密
aes_key = b'APP-Cumminschine'
iv_key = b'Cty-china-ER-app'
encrypt_content = b'python_Cummins'

imgs=[]
vins=[]
esns=[]
dates=[]
model_rd=[] #道路
model_off=[] #非道路
comp=[]

# 图片路径
# image_file_path = "/tmp/"
# model_path = r"models/model_draft_1129.pth"
image_file_path = "/Users/ctysmac/Desktop/vsCode/ocr_local_test/src/image/"
model_path = r"models/model_draft_1129.pth"

loaded_model = None
ocr = PaddleOCR(
    use_angle_cls=True,
    cls_model_dir=r"models/cls",
    det_model_dir=r"models/det",
    rec_model_dir=r"models/rec"
)
# 角度不对 模糊 假的  未识别
class_labels = ['您拍摄照片的角度倾斜', '您拍摄的照片比较模糊', '您拍摄的照片可能存在翻拍', '']


# 上传文件
# type (1  ESN) (2 OEMName、ChassisNumber、VIN) (3 OEMName、VIN) (4、WarrantyStartDate、ChassisNumber) (5、WarrantyStartDate)
def upload(file_content, file_name, ocr_type):
    file_path = image_file_path + file_name
    logger.info(f"saving file into : {file_path}")
    # __save_file(file_path, file_content)

    ocr_result = run_model(file_path,ocr_type)

    msg = ""
    if not ocr_result:
        msg = __single_cf_test(file_path)

    ocr_results = __dataIntegration(ocr_result,file_path,ocr_type)

    logger.info("识别内容" + str(ocr_results) + "-------msg:" + msg)

    # __clear_file(file_path)

    return {"code": 1, "data": ocr_results, 'msg': msg}

#数据整合
def __dataIntegration(ocr_result,file_path,ocr_type):
    ocr_results = {}
    if ocr_result and ocr_type in [2, 3]:
        if ocr_result[0] != 'no':
            ocr_results["VIN"] = ocr_result[0]

    if ocr_result and ocr_type == 1:
        if ocr_result[1] != 'no':
            ocr_results["ESN"] = ocr_result[1]

    if ocr_result and ocr_type in [4, 5]:
        if ocr_result[2] != 'no':
            ocr_results["WarrantyStartDate"] = ocr_result[2]

    if ocr_result and ocr_type == 4:
        if ocr_result[3] != 'no':
            ocr_results["ChassisNumber"] = ocr_result[3]

    if ocr_result and ocr_type == 2:
        if ocr_result[4] != 'no':
            ocr_results["ChassisNumber"] = ocr_result[4]

    if ocr_result and ocr_type in [2, 3]:
        if ocr_result[5] != 'no':
            ocr_results["OEMName"] = ocr_result[5]
    return ocr_results
def run_model(img_path,ocr_type):
    # =========================================================
    # imgpil = Image.open(img_path)
    # width, height = imgpil.size
    # new_height = int(0.85 * height)
    # cropped_img = imgpil.crop((0, 0, width, new_height))
    # # cropped_img=cropped_img.rotate(90)
    # cropped_img.save('./temp_img.jpg')
    # img_path='./temp_img.jpg'
    # =========================================================

    recgs = []
    for i in range(0, 5):

        if i == 0:

            result = ocr.ocr(img_path)
            if result != [None]:
                txts = [line[1][0] for line in result[0]]
            else:
                txts = ['0']
            ratios = []
            boxes = [line[0] for line in result[0]]
            i = 0
            for word_info in boxes:
                # print(txts[i])
                i += 1
                width = word_info[1][0] - word_info[0][0]
                height = word_info[2][1] - word_info[1][1]
                ratio = width / height
                # print(word_info)
                # print(ratio)
                ratios.append(ratio)
            average_ratio = np.mean(ratios)

        if i == 1:
            pre_img = cv2.imread(img_path)
            h, w = pre_img.shape[:2]
            if h > w:
                pre_img = cv2.rotate(pre_img, cv2.ROTATE_90_COUNTERCLOCKWISE)
            if max(h, w) > 1080:
                pre_img = img_resize(pre_img)
            result = ocr.ocr(pre_img)

        if i == 2:
            pre_img = cv2.imread(img_path)
            h, w = pre_img.shape[:2]
            if h > w:
                pre_img = cv2.rotate(pre_img, cv2.ROTATE_90_CLOCKWISE)
            if max(h, w) > 1080:
                pre_img = img_resize(pre_img)
            result = ocr.ocr(pre_img)

        if i == 3:
            pre_img = cv2.imread(img_path)
            h, w = pre_img.shape[:2]
            pre_img = cv2.rotate(pre_img, cv2.ROTATE_180)
            if max(h, w) > 1080:
                pre_img = img_resize(pre_img)

            result = ocr.ocr(pre_img)

        if i == 4:
            pre_img = cv2.imread(img_path)
            # h,w = pre_img.shape[:2]
            pre_img = cv2.rotate(pre_img, cv2.ROTATE_90_COUNTERCLOCKWISE)
            # if max(h,w)>1080:
            #     pre_img = img_resize(pre_img)

            result = ocr.ocr(pre_img)

        if result != [None]:
            txts = [line[1][0] for line in result[0]]
        else:
            txts = ['0']

        scores = [line[1][1] for line in result[0]]
        recgs.append(recognition_process(txts, scores,ocr_type))
        # rt.append(time.time())
        # 初始化变量来存储最小计数和对应的子数组
    # print(recgs)
    best_result = result_select(recgs)
    print(best_result)
    return best_result

def img_resize(pre_img):
    h,w = pre_img.shape[:2]
    ratio =1080/w
    pre_img = cv2.resize(pre_img, (1080,int(h*ratio)))
    return pre_img

def result_select(recgs):
    min_no_count = float('inf')
    best_result = None
    for arr in recgs:
        count_no = arr.count('no')
        if count_no <= min_no_count:
            min_no_count = count_no
            best_result = arr.copy()

    for i, value in enumerate(best_result):
        longest_value = value
        for other_arr in recgs:
            other_value = other_arr[i]
            if len(other_value) > len(longest_value) and other_value != 'no':
                longest_value = other_value
        best_result[i] = longest_value
    return best_result


def esn_relate(arr):
    lower_arr = [item.lower() for item in arr]

    esn_index = None
    for i, item in enumerate(lower_arr):
        if item == 'esn':
            esn_index = i
            break

    if esn_index is None:

        for item in arr:
            if item.isdigit() and len(item) == 8:
                return item
        return None

    min_distance = float('inf')
    closest_number = None

    for i, item in enumerate(arr):
        if item.isdigit() and len(item) == 8:
            distance = abs(i - esn_index)
            if distance < min_distance:
                min_distance = distance
                closest_number = item

    return closest_number


def merge_strings_within_distance(input_list):
    result = []
    pattern_f2 = re.compile(r'^F2\.?$')
    pattern_8char = re.compile(r'^\.?(?=.*[A-Za-z])(?=.*\d)[A-Za-z\d]{8}$')

    i = 0
    while i < len(input_list):
        current = input_list[i]

        if pattern_f2.match(current):
            merged = current

            for j in range(1, 4):
                if i + j < len(input_list) and pattern_8char.match(input_list[i + j]):
                    merged = f'F2.{input_list[i + j].lstrip(".")}'
                    i += j
                    break
            result.append(merged)
        else:
            if pattern_8char.match(current):
                merged = None
                for k in range(1, 4):
                    if i - k >= 0 and pattern_f2.match(input_list[i - k]):
                        merged = f'F2.{current.lstrip(".")}'
                        result[-k] = merged
                        break
                    elif i + k < len(input_list) and pattern_f2.match(input_list[i + k]):
                        merged = f'F2.{current.lstrip(".")}'
                        i += k
                        break
                if merged:
                    result.append(merged)
                else:
                    result.append(current)
            else:
                result.append(current)

        i += 1

    return result

def recognition_process(txts, scores, ocr_type):
    vincount = 0
    datecount = 0
    esncount = 0
    rdcount = 0
    offrdcount = 0
    compcount = 0

    if esn_relate(txts) != None:
        esns.append(esn_relate(txts))
        esncount = 1
    # print(txts)
    txts = merge_strings_within_distance(txts)
    # print('###############')
    # print(txts)
    for i in range(len(txts)):
        # print(txts[i],'==',scores[i])
        v_vin = re.findall("[XL]{1}[A-Z0-9]{16}", str(txts[i]))
        if (v_vin != []) & (vincount == 0):
            if any(char.isdigit() for char in str(txts[i])):
                v_vin[0] = v_vin[0].replace('O', '0')
                vins.append(v_vin[0])
                vincount = 1
        # 移除
        pattern = r'(?<!\d)(\d{8})(?!\d)'
        match = re.search(pattern, str(txts[i]))
        if match:
            if (esncount == 0):
                esns.append(match.group(1))
                esncount = 1

        v_date = re.findall("\d{4}-\d{2}-\d{2}", str(txts[i]))
        if (v_date != []) & (datecount == 0):
            dates.append(v_date[0])
            datecount = 1
        if (v_date != []) & (datecount != 0):
            try:
                v_date[0] = v_date[0].replace('88', '08')
                dates[-1] = dates[-1].replace('88', '08')
                date_a = datetime.strptime(v_date[0], '%Y-%m-%d')
                date_b = datetime.strptime(dates[-1], '%Y-%m-%d')
                if date_a > date_b:
                    dates[-1] = v_date[0]
            except:

                a = 0

        v_model_rd = re.findall(r'[\u4e00-\u9fa5]+牌[A-Z0-9]+.*', str(txts[i]))
        if (v_model_rd != []) & (rdcount == 0):
            # print('model_rd',v_model_rd[0])
            v_model_rd[0] = v_model_rd[0].replace('品牌型号', '')
            chinese_pattern = re.compile(r'[\u4e00-\u9fff]+')
            if '放' in v_model_rd[0]:
                v_model_rd[0] = chinese_pattern.sub('解放牌', v_model_rd[0])
            model_rd.append(v_model_rd[0])
            rdcount = 1

        v_model_off = re.findall(r'^F[0-9]+\.[A-Z0-9]+$', str(txts[i]))
        if (v_model_off != []) & (offrdcount == 0):
            # print('model_off',v_model_off[0])
            model_off.append(v_model_off[0])
            offrdcount = 1
        v_model_off2 = re.findall(r'F2\.[A-Za-z0-9]{8}', str(txts[i]))
        if (v_model_off2 != []) & (offrdcount == 0):
            # print('model_off',v_model_off[0])
            model_off.append(v_model_off2[0])
            offrdcount = 1

        v_model_off3 = re.findall(r'[X](?=.*[0-9])[0-9a-zA-Z]{5,10}(?<![kK][wW]|kg)$', str(txts[i]))
        if (v_model_off3 != []) & (offrdcount == 0):
            # print('model_off',v_model_off[0])
            model_off.append(v_model_off3[0])
            offrdcount = 1

        v_comp = re.findall(r'.*有限.*', str(txts[i]))
        if (v_comp != []) & (compcount == 0):
            v_comp[0] = v_comp[0].replace('中国', '')
            v_comp[0] = v_comp[0].replace('制造', '')
            # print('comp',v_comp[0])
            comp.append(v_comp[0])
            compcount = 1

    if vincount == 0:
        vins.append('no')
    if datecount == 0:
        dates.append('no')
    if esncount == 0:
        esns.append('no')
    if rdcount == 0:
        model_rd.append('no')
    if offrdcount == 0:
        model_off.append('no')
    if compcount == 0:
        comp.append('no')
    else:
        oem_df = pd.read_csv('data/oem_list_from_sf_claims.csv')
        oem_list = oem_df['OEM名称'].values

        # print(comp)
        closest_match = find_closest_match(comp[-1], oem_list)
        if closest_match:
            comp[-1] = closest_match
    return [vins[-1],esns[-1],dates[-1],model_rd[-1],model_off[-1],comp[-1]]

    # ocr_results = {}
    # if vins and ocr_type in [2, 3]:
    #     ocr_results["VIN"] = vins[0]
    #
    # if esns and ocr_type == 1:
    #     ocr_results["ESN"] = esns[0]
    #
    # if dates and ocr_type in [4, 5]:
    #     ocr_results["WarrantyStartDate"] = dates[0]
    #
    # if model_rd and ocr_type == 4:
    #     ocr_results["ChassisNumber"] = model_rd[0]
    #
    # if model_off and ocr_type == 2:
    #     ocr_results["ChassisNumber"] = model_off[0]
    #
    # if comp and ocr_type in [2, 3]:
    #     ocr_results["OEMName"] = comp[0]
    # return ocr_results

# 识别失败-报错处理
def __single_cf_test(image_path):
    global loaded_model
    image = Image.open(image_path).convert('RGB')
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    input_tensor = transform(image)
    input_batch = input_tensor.unsqueeze(0)
    if loaded_model is None:
        loaded_model = torch.load(model_path, map_location=torch.device('cpu'))
    loaded_model.eval()
    with torch.no_grad():
        output = loaded_model(input_batch)

    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    predicted_class = torch.argmax(probabilities).item()
    predicted_label = class_labels[predicted_class]

    return predicted_label

def levenshtein_distance(s1, s2):
    if len(s1) < len(s2):
        return levenshtein_distance(s2, s1)

    if len(s2) == 0:
        return len(s1)

    previous_row = range(len(s2) + 1)

    for i, c1 in enumerate(s1):
        current_row = [i + 1]

        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)

            current_row.append(min(insertions, deletions, substitutions))

        previous_row = current_row

    return previous_row[-1]

def find_closest_match(oem_name, name_list):
    best_match = None
    lowest_distance = float('inf')

    for name in name_list:
        distance = levenshtein_distance(oem_name, name)
        if distance < lowest_distance:
            lowest_distance = distance
            best_match = name
    if lowest_distance/len(best_match)<0.6:
        return best_match

# 加密
def encrypt(plaintext):
    encryptor = Cipher(
        algorithms.AES(aes_key),
        modes.GCM(iv_key),
        backend=default_backend()
    ).encryptor()
    # 将tag直接追加在最后，即可和java解密代码兼容
    ciphertext = encryptor.update(plaintext) + encryptor.finalize() + encryptor.tag
    return ciphertext.hex().upper()


# 解密
def decrypt(plaintext):
    dec = Cipher(
        algorithms.AES(aes_key),
        modes.GCM(iv_key),
        backend=default_backend()
    ).decryptor()
    # 将tag直接追加在最后，即可和java解密代码兼容
    ciphertext = dec.update(plaintext) + dec.finalize() + dec.tag
    return ciphertext.hex().upper()


# save file to local path
def __save_file(file_path, file_content):
    try:
        with open(file_path, "wb") as lf:
            lf.write(file_content)
        return True
    except Exception as e:
        logger.error('成功保存失败发生错误：', str(e))
        return False


def __clear_file(file_path):
    try:
        os.remove(file_path)
    except Exception as e:
        logger.error('clear file error', str(e))


def return_result(code=200, data={}):
    return {
        "statusCode": code,
        "headers": {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*',
            'Access-Control-Allow-Headers': '*',
            'Access-Control-Allow-Methods': '*'
        },
        "body": json.dumps(data)
    }


def lambda_handler(event, context):
    try:
        headers = event.get("headers")
        body = event.get('body')
        ocr_type = headers.get("type")
        token = headers.get("token")
        file_name = headers.get("filename")

        if ocr_type is None:
            raise ValueError("type is invalid")
        if body is None:
            raise ValueError("body is invalid")
        if file_name is None:
            raise ValueError("filename is invalid")
        if token is None:
            raise ValueError("token is invalid")

        ciphertext = encrypt(encrypt_content)
        if ciphertext != token:
            logger.error('密钥不对')
            return return_result(data={"code": 0, "msg": "密钥不对"})

        file_content = base64.b64decode(body.encode('utf-8'))

        result = upload(file_content, file_name, int(ocr_type))

        return return_result(data=result)
    except:
        return return_result(data={"code": 0, "msg": ""})

upload("/Users/ctysmac/Desktop/ocr_local_test/src/image/file_name.jpg","file_name.jpg",2)
# upload("/Users/ctysmac/Desktop/vsCode/ocr_local_test/OCR_update_0528/002_OCR_Detector_update_2405_v4_speed_0528.py","file_name",2)
import paddle
import sys
print(paddle.__version__)
print(sys.version)